import os
import sys
import logging
import time
import numpy as np
from onnxruntime import InferenceSession
from transformers import T5Tokenizer

from .. import helper
from ..onnx_base import ONNXBase


def update_eosflag(eos_flag, is_eos):
    eos_flag = np.multiply(eos_flag, (~is_eos).astype(np.int32))

    return eos_flag


class ONNXT5DecInfer(ONNXBase):
    def __init__(self, *args) -> None:
        super(ONNXT5DecInfer, self).__init__(*args)
        logging.debug(f"ONNXT5DecInfer init.")
        self.max_dec_num = 32

        dec_onnx_path = f"{self.model_dir}/decoder.onnx"
        self.dec_sess = InferenceSession(dec_onnx_path, providers=self.providers)
        self.dec_inputs_name = self._get_onnx_input_name(self.dec_sess)
        self.dec_outputs_name = self._get_onnx_output_name(self.dec_sess)

        lm_onnx_path = f"{self.model_dir}/lm_head.onnx"
        self.lm_sess = InferenceSession(lm_onnx_path, providers=self.providers)
        self.lm_inputs_name = self._get_onnx_input_name(self.lm_sess)
        self.lm_outputs_name = self._get_onnx_output_name(self.lm_sess)

    def __del__(self):
        super(ONNXT5DecInfer, self).__del__()
        logging.debug(f"ONNXT5DecInfer destruct.")

    def onnx_infer(self, idx, dec_ids, enc_mask, enc_hidden_state):
        start_time = time.time()
        dec_hidden_state, = self.dec_sess.run(self.dec_outputs_name, 
                                            {self.dec_inputs_name[0]:dec_ids, 
                                            self.dec_inputs_name[1]:enc_mask,
                                            self.dec_inputs_name[2]:enc_hidden_state})

        dec_time = (time.time() - start_time) * 1000
        start_time = time.time()
        logits, = self.lm_sess.run(self.lm_outputs_name, 
                                {self.lm_inputs_name[0]:dec_hidden_state[:, idx, :]})  # 注意优化点!!!
        lm_time = (time.time() - start_time) * 1000
        logging.debug(f"shape decoder:{dec_hidden_state.shape} lm:{logits.shape}")
        logging.debug(f"time decoder:{dec_time:.3f} (ms) lm_head:{lm_time:.3f} (ms)")

        return logits

    def dec(self, enc_mask, enc_hidden_state):
        eos_token_id = 1
        pad_token_id = 0
        eos_flag = np.ones(enc_mask.shape[0]).astype(np.int32)
        dec_ids = np.zeros((1, self.max_dec_num), dtype=np.int32)

        dec_start = time.time()
        for i in range(self.max_dec_num):
            logits = self.onnx_infer(i, dec_ids, enc_mask, enc_hidden_state)
            next_tokens = np.argmax(logits, axis=-1)
            logging.debug(f"np argmax next_tokens:{next_tokens}")
            ### base value: 21126, 43, 2008, 24, 293, 53, 3, 9, 1782, 19, 207, 21, 25, 3, 5, 1

            next_tokens = next_tokens * eos_flag + pad_token_id * (1 - eos_flag)
            dec_ids[:, i+1] = next_tokens

            eos_flag = update_eosflag(eos_flag, next_tokens == eos_token_id)
            if eos_flag.max() == 0:
                break

        dec_time = (time.time() - dec_start) * 1000
        avg_time = dec_time / (i+1)
        logging.info(f"dec num:{i+1} dec_all:{dec_time:.3f} (ms) dec_avg: {avg_time:.3f} (ms)")
        logging.debug(f"dec result:{dec_ids}")

        return dec_ids